import tkinter as tk
from tkinter import scrolledtext, Frame, Button, Entry, Label, messagebox, filedialog
import threading
import os
import time
import json
import re
import glob
import base64
import requests
from dotenv import load_dotenv
from PIL import Image, ImageTk
from io import BytesIO
import datetime
import random
import requests
from openai import OpenAI
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from webdriver_manager.microsoft import EdgeChromiumDriverManager
from selenium.webdriver.edge.service import Service as EdgeService
from selenium.webdriver.edge.options import Options as EdgeOptions
import logging
import sys
import traceback
import signal
import numpy as np
import json
import os#文件系统，进程，环境变量，posix的封装
import time#time低级操作
import glob#通配符
from datetime import datetime#time类型的高级操作
from openai import OpenAI
from dotenv import load_dotenv#.env文件的使用
import requests#crawl爬取
from bs4 import BeautifulSoup#HTML/XML解析与数据提取
from selenium import webdriver#浏览器自动化控制
from selenium.webdriver.common.by import By#提供统一的元素定位标识符，用于指定查找元素的方式（如 ID、XPath、CSS 等）
from selenium.webdriver.support.ui import WebDriverWait#智能等待页面元素加载
from selenium.webdriver.support import expected_conditions as EC#定义页面加载完成的条件
from selenium.webdriver.edge.service import Service as EdgeService#管理 Edge 浏览器驱动（msedgedriver.exe）的生命周期，包括启动/停止服务
from selenium.webdriver.edge.options import Options as EdgeOptions#定制浏览器启动参数（如无头模式、代理、证书设置）
import re#正则表达式文本匹配
import logging#日志记录与管理
import random#随机数
import lxml#高性能HTML/XML解析器
import numpy as np#数值计算和数组操作
from typing import Optional
from pathlib import Path
import os
import base64
import requests
from PIL import Image
from io import BytesIO
import os
from pathlib import Path
from typing import Dict, List, Optional
from pathlib import Path
from typing import Dict, List, Optional
import datetime as dt 
from tkinter import ttk
from PIL import Image, ImageEnhance
import tempfile
from tkinter import simpledialog
# 加载环境变量
load_dotenv()

# 常量路径的定义
HISTORY_DIR = "chat_histories"
CRAWLER_DIR = "crawled_data"
KNOWLEDGE_DIR = CRAWLER_DIR
KEYWORDS = ["咨询", "热线", "电话", "联系", "联系方式", "联系电话", "招生热线"]
# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',#定义日志输出的格式字符串
    handlers=[
        logging.FileHandler('chat_system.log'),#将日志写入文件chat_system.log
        logging.StreamHandler()#将日志输出到控制台（如终端）
    ]
)
logger = logging.getLogger(__name__)

# 解决 datetime 问题的方法：统一使用 datetime.datetime
# 创建辅助函数来获取当前时间
def get_current_timestamp():
    """获取当前时间戳"""
    return dt.datetime.now().strftime("%Y%m%d_%H%M%S")

def get_current_datetime():
    """获取当前日期时间"""
    return dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# 在代码中使用这些辅助函数替代直接调用 datetime.datetime.now()
# ======================== 核心功能模块 - 与命令行版本完全一致 ========================

class WebCrawler:
    """网页爬取核心类"""
    def __init__(self, output_dir=CRAWLER_DIR):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.driver = None
        self.user_agents = [
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36...",
            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)...",
            "Mozilla/5.0 (iPhone; CPU iPhone OS 16_0 like Mac OS X)..."
        ]

    def _init_driver(self):
        if self.driver is None:
            edge_options = EdgeOptions()
            edge_options.add_argument('--headless=new')
            edge_options.add_argument('--disable-gpu')
            edge_options.add_argument('--log-level=3')
            edge_options.add_argument('--ignore-certificate-errors')
            edge_options.add_argument('--allow-running-insecure-content')
            edge_options.add_argument('--disable-blink-features=AutomationControlled')
            edge_options.add_experimental_option("excludeSwitches", ["enable-automation"])
            
            edge_options.add_argument(f'--user-agent={random.choice(self.user_agents)}')
            edge_options.add_argument('--window-size=1920,1080')
            
            service = EdgeService(
                EdgeChromiumDriverManager().install(),
                service_args=["--silent"],
                log_path=os.devnull
            )
            self.driver = webdriver.Edge(service=service, options=edge_options)
        return True

    def _get_dynamic_content(self, url, timeout=20):
        if not self._init_driver():
            return None
            
        try:
            self.driver.get(url)
            
            WebDriverWait(self.driver, timeout).until(
                lambda d: d.execute_script("return document.readyState") == "complete" and
                         len(d.find_elements(By.TAG_NAME, 'body')) > 0
            )
            
            for i in range(3):
                scroll_height = self.driver.execute_script(
                    "return Math.min(document.body.scrollHeight, "
                    f"document.documentElement.scrollHeight)")
                self.driver.execute_script(
                    f"window.scrollTo(0, {scroll_height * (i+1)/3})")
                time.sleep(random.uniform(0.5, 1.5))
            
            WebDriverWait(self.driver, 5).until(
                EC.presence_of_element_located((By.XPATH, "//*[text()]"))
            )
            
            return self.driver.page_source
            
        except Exception as e:
            try:
                return requests.get(url, verify=False, timeout=10).text
            except:
                return None
        finally:
            if self.driver:
                self.driver.quit()
                self.driver = None
    
    def _clean_text(self, text):
        cleaned = re.sub(r'\s+', ' ', text).strip()
        return re.sub(r'[\x00-\x1f\x7f-\x9f]', '', cleaned)
    
    def crawl_page(self, url, is_dynamic=False, selector=None):
        try:
            if is_dynamic:
                html = self._get_dynamic_content(url)
            else:
                response = requests.get(url, headers={
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                }, timeout=10)
                response.encoding = "utf-8"
                html = response.text if response.status_code == 200 else None
            
            if not html:
                return {"status": "failed", "reason": "页面获取失败"}
            
            soup = BeautifulSoup(html, 'html.parser')
            
            if selector:
                if selector.startswith('/'):
                    from lxml import html
                    tree = html.fromstring(html)
                    elements = tree.xpath(selector)
                    content = '\n'.join([self._clean_text(e.text_content()) for e in elements if e.text_content()])
                else:
                    elements = soup.select(selector)
                    content = '\n'.join([self._clean_text(e.get_text()) for e in elements if e.get_text()])
            else:
                content = self._clean_text(soup.get_text())
            
            # 使用辅助函数获取时间戳
            timestamp = get_current_timestamp()
            domain = re.sub(r'[^\w-]', '_', url.split('//')[-1].split('/')[0])
            filename = f"{domain}_{timestamp}.txt"
            filepath = os.path.join(self.output_dir, filename)
            
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(f"URL: {url}\n\n{content}")
            
            return {
                "status": "success",
                "path": filepath,
                "content": content[:1000] + "..." if len(content) > 1000 else content
            }
            
        except Exception as e:
            return {"status": "error", "reason": str(e)}

class HistoryManager:
    """历史记录管理类"""
    def __init__(self, history_dir=HISTORY_DIR):
        self.history_dir = history_dir
        os.makedirs(history_dir, exist_ok=True)
    
    def list_histories(self, limit=10):
        files = []
        for fname in os.listdir(self.history_dir):
            if fname.endswith(".json") and fname != "latest.json":
                file_path = os.path.join(self.history_dir, fname)
                files.append({
                    "id": len(files) + 1,
                    "filename": fname,
                    "path": file_path,
                    "mtime": os.path.getmtime(file_path)
                })
        
        files.sort(key=lambda x: x["mtime"], reverse=True)
        return files[:limit]
    
    def load_history(self, filename):
        file_path = os.path.join(self.history_dir, filename)
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"加载历史记录失败: {filename} - {e}")
        return None
    
    def load_multiple_histories(self, filenames):
        merged_history = []
        for filename in filenames:
            history = self.load_history(filename)
            if history:
                if merged_history and history[0].get("role") == "system":
                    merged_history.extend(history[1:])
                else:
                    merged_history.extend(history)
        return merged_history
    
    def save_history(self, messages, custom_name=None):
        timestamp = get_current_timestamp()
        filename = custom_name or f"history_{timestamp}.json"
        filepath = os.path.join(self.history_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        latest_path = os.path.join(self.history_dir, "latest.json")
        with open(latest_path, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        return filepath

class KnowledgeLoader:
    """知识库加载类"""
    def __init__(self, knowledge_dir=KNOWLEDGE_DIR):
        self.knowledge_dir = knowledge_dir
        os.makedirs(knowledge_dir, exist_ok=True)
    
    def load(self):
        content = []
        extensions = ('.txt', '.md', '.csv', '.json')
        
        for ext in extensions:
            pattern = os.path.join(self.knowledge_dir, '**', f'*{ext}')
            for file_path in glob.glob(pattern, recursive=True):
                if os.path.isfile(file_path):
                    try:
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            text = f.read(10240) if os.path.getsize(file_path) > 102400 else f.read()
                            content.append({
                                "path": os.path.basename(file_path),
                                "content": text
                            })
                    except Exception as e:
                        print(f"文件读取失败: {file_path} - {e}")
        return content
    
    def kload(self, file_filter=None, content_filter=None):
        content = []
        extensions = ('.txt', '.md', '.csv', '.json')
        project_root = os.path.dirname(self.knowledge_dir)
        search_dirs = [project_root, os.path.join(project_root, "models"), self.knowledge_dir]
        
        if isinstance(file_filter, str) and (os.path.isabs(file_filter) or os.path.exists(file_filter)):
            file_paths = [file_filter]
        else:
            file_paths = []
            if isinstance(file_filter, list):
                for fname in file_filter:
                    for base_dir in search_dirs:
                        path = os.path.join(base_dir, fname)
                        if os.path.isfile(path):
                            file_paths.append(path)
                            break
            elif file_filter:
                pattern = f"*{file_filter}*"
                for base_dir in search_dirs:
                    if os.path.exists(base_dir):
                        search_pattern = os.path.join(base_dir, pattern)
                        found_files = glob.glob(search_pattern, recursive=True)
                        file_paths.extend(found_files)
            else:
                for base_dir in search_dirs:
                    if os.path.exists(base_dir):
                        for ext in extensions:
                            pattern = os.path.join(base_dir, '**', f'*{ext}')
                            files = glob.glob(pattern, recursive=True)
                            file_paths.extend(files)
        
        validated_paths = []
        for path in file_paths:
            try:
                if os.path.exists(path) and os.path.isfile(path):
                    validated_paths.append(path)
            except:
                pass
        
        for file_path in validated_paths:
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    file_size = os.path.getsize(file_path)
                    text = f.read(10240) if file_size > 102400 else f.read()
                    
                    if content_filter and not content_filter(text):
                        continue
                    
                    content.append({
                        "path": file_path,
                        "filename": os.path.basename(file_path),
                        "content": text
                    })
            except Exception as e:
                print(f"文件读取失败: {file_path} - {e}")
        
        return content
    def save_history(self, messages, custom_name=None):
        """保存历史记录 - 修复 datetime 问题"""
        # 使用统一的 get_current_timestamp 函数
        # 使用辅助函数获取时间戳
        timestamp = get_current_timestamp()
        filename = custom_name or f"history_{timestamp}.json"
        filepath = os.path.join(self.history_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        latest_path = os.path.join(self.history_dir, "latest.json")
        with open(latest_path, 'w', encoding='utf-8') as f:
            json.dump(messages, f, ensure_ascii=False, indent=4)
        
        return filepath

class ImageOCR:
    """使用 GLM-4V 模型进行图片 OCR 识别 - 优化版本"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
    
    def preprocess_image(self, image_path: str) -> str:
        """图片预处理 - 提高OCR准确率"""
        try:
            # 打开图片
            img = Image.open(image_path)
            
            # 转换为灰度图
            img = img.convert('L')
            
            # 增强对比度
            enhancer = ImageEnhance.Contrast(img)
            img = enhancer.enhance(1.5)  # 1.5倍对比度增强
            
            # 创建临时文件
            temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
            temp_path = temp_file.name
            img.save(temp_path, "JPEG", quality=95)
            
            return temp_path
        except Exception as e:
            print(f"图片预处理失败: {str(e)}")
            return image_path  # 失败时返回原图
    def verify_api_key(self):
        """验证API密钥有效性"""
        test_payload = {
            "model": "glm-4v",
            "messages": [{"role": "user", "content": "这是一次API密钥验证请求"}],
            "max_tokens": 5
        }
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        try:
            response = requests.post(
                self.base_url, 
                headers=headers, 
                json=test_payload,
                timeout=10
            )
            
            if response.status_code == 200:
                print("✅ API密钥验证通过")
                return True
            elif response.status_code == 401:
                raise ValueError("API密钥无效或已过期")
            elif response.status_code == 403:
                raise ValueError("API密钥权限不足或账户欠费")
            else:
                raise ValueError(f"API请求失败: {response.status_code} - {response.text}")
        except requests.exceptions.ConnectionError:
            raise ValueError("无法连接到API服务器，请检查网络连接")
        except requests.exceptions.Timeout:
            raise ValueError("API请求超时，请稍后重试")
        except Exception as e:
            raise ValueError(f"API密钥验证失败: {str(e)}")
    def recognize_image(self, image_path: str, max_retries=3) -> Optional[str]:
        """
        使用GLM-4V模型识别图片中的文字 - 优化版本
        
        Args:
            image_path (str): 图片文件路径
            max_retries (int): 最大重试次数
            
        Returns:
            Optional[str]: 识别出的文字内容，失败时返回None
        """
        # 预处理图片
        processed_path = self.preprocess_image(image_path)
        
        try:
            # 读取图片并转换为base64
            with open(processed_path, "rb") as image_file:
                base64_image = base64.b64encode(image_file.read()).decode('utf-8')
            
            # 构建请求头
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "application/json"
            }
            
            # 构建请求体 - 优化提示词
            payload = {
                "model": "glm-4v",
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text", 
                                "text": (
                                    "请仔细识别图片中的所有文字内容，包括印刷体和手写体。"
                                    "保持原始格式和排版，特别注意表格、图表中的文字。"
                                    "对于模糊或难以识别的文字，请尽力猜测并标记为[?]。"
                                    "输出结果时请确保准确性，不要遗漏任何文字。"
                                )
                            },
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
                            }
                        ]
                    }
                ],
                "max_tokens": 3000,  # 增加token限制
                "temperature": 0.2   # 降低随机性
            }
            
            # 重试机制
            for attempt in range(max_retries):
                try:
                    # 发送请求
                    response = requests.post(
                        self.base_url, 
                        headers=headers, 
                        json=payload,
                        timeout=30  # 30秒超时
                    )
                    
                    # 检查响应状态
                    if response.status_code == 200:
                        # 解析响应
                        result = response.json()
                        if 'choices' in result and len(result['choices']) > 0:
                            content = result['choices'][0]['message']['content']
                            return content
                        else:
                            print(f"API响应格式错误 (尝试 {attempt+1}/{max_retries})")
                            print(f"完整响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
                    else:
                        print(f"API请求失败 (尝试 {attempt+1}/{max_retries}): {response.status_code}")
                        print(f"错误信息: {response.text}")
                    
                    # 等待后重试
                    time.sleep(1)
                    
                except requests.exceptions.Timeout:
                    print(f"API请求超时 (尝试 {attempt+1}/{max_retries})")
                    time.sleep(2)
                except requests.exceptions.RequestException as e:
                    print(f"网络请求错误 (尝试 {attempt+1}/{max_retries}): {str(e)}")
                    time.sleep(1)
                except Exception as e:
                    print(f"OCR处理出错 (尝试 {attempt+1}/{max_retries}): {str(e)}")
                    time.sleep(1)
            
            return None
            
        except Exception as e:
            print(f"OCR处理出错: {str(e)}")
            return None
        finally:
            # 清理临时文件
            if processed_path != image_path and os.path.exists(processed_path):
                os.remove(processed_path)
    
    def save_ocr_result(self, image_path: str, output_dir: str = "knowledge") -> Optional[str]:
        """
        识别图片文字并保存为文本文件
        
        Args:
            image_path (str): 图片文件路径
            output_dir (str): 输出目录，默认为"knowledge"
            
        Returns:
            Optional[str]: 保存的文件路径，失败时返回None
        """
        # 识别图片文字
        text_content = self.recognize_image(image_path)
        if not text_content:
            print("OCR识别失败")
            return None
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成输出文件名
        image_name = os.path.basename(image_path)
        txt_name = os.path.splitext(image_name)[0] + ".txt"
        output_path = os.path.join(output_dir, txt_name)
        
        # 保存文本文件
        try:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(text_content)
            print(f"✅ OCR结果已保存至: {output_path}")
            return output_path
        except Exception as e:
            print(f"保存OCR结果失败: {str(e)}")
            return None
# 加载环境变量
load_dotenv()

# 初始化OpenAI客户端
client = OpenAI(
    api_key="sk-uZy29fuiHwjQEHtJq35r5KQHDqYYMa3WNdVdvxo8ZoPyws4T",
    base_url="http://123.182.124.194:58000/v1"
)
###################################embedding#####################################################
# 预定义关键词库（可根据需要扩展）

# 预定义关键词库（可根据需要扩展）
KEYWORDS = ["咨询", "热线", "电话", "联系", "联系方式", "联系电话", "招生热线"]

def load_documents(data_dir, max_length=500):
    """从指定目录加载文档，智能分段处理"""
    documents = []
    files = glob.glob(os.path.join(data_dir, '*.txt'))
    
    if not files:
        print(f"⚠️ 在 '{data_dir}' 目录中未找到任何txt文件")
        return documents
    
    print(f"📂 发现 {len(files)} 个文本文件")
    
    for file_path in files:
        try:
            with open(file_path, "r", encoding='utf-8', errors='ignore') as f:
                content = f.read().strip()
                if not content:
                    continue
                    
                filename = os.path.basename(file_path)
                
                # 短文档直接整体处理
                if len(content) <= max_length:
                    has_keywords = any(kw in content for kw in KEYWORDS)
                    documents.append({
                        "source": filename,
                        "content": content,
                        "has_keywords": has_keywords
                    })
                # 长文档智能分段
                else:
                    segments = smart_segment(content, max_length)
                    if not segments:
                        continue
                        
                    # 标识包含关键词的片段
                    for i, seg in enumerate(segments):
                        has_keywords = any(kw in seg for kw in KEYWORDS)
                        documents.append({
                            "source": f"{filename}@{i+1}",
                            "content": seg,
                            "has_keywords": has_keywords
                        })
        except Exception as e:
            print(f"⚠️ 加载文档出错: {os.path.basename(file_path)} - {str(e)}")
    
    # 统计包含关键词的文档数
    key_docs = sum(1 for d in documents if d['has_keywords'])
    print(f"✅ 加载完成: {len(documents)}个文本片段 ({key_docs}个包含关键词)")
    return documents

def smart_segment(text, max_length=500):
    """智能文本分段：优先保留含关键词的上下文"""
    segments = []
    
    # 首先查找关键词周围的上下文
    found_key_contexts = []
    for kw in KEYWORDS:
        if kw in text:
            start_idx = max(0, text.index(kw) - 50)
            end_idx = min(len(text), text.index(kw) + 150)
            segment = text[start_idx:end_idx]
            if segment not in found_key_contexts:
                found_key_contexts.append(segment)
    
    # 添加找到的关键词上下文
    segments.extend(found_key_contexts)
    
    # 剩余文本按自然段落分割
    remaining_text = text
    for seg in segments:
        remaining_text = remaining_text.replace(seg, "")
    
    if len(remaining_text) > max_length:
        # 按段落分割
        if '\n\n' in remaining_text:
            paras = remaining_text.split('\n\n')
            for para in paras:
                if para and len(para) <= max_length:
                    segments.append(para)
                elif para:
                    # 按句子分割
                    sentences = re.split(r'(?<=[。！？.!?])', para)
                    current = ""
                    for sent in sentences:
                        if len(current) + len(sent) > max_length:
                            segments.append(current)
                            current = sent
                        else:
                            current += sent
                    if current:
                        segments.append(current)
        else:
            # 简单长度分割
            while remaining_text:
                segments.append(remaining_text[:max_length])
                remaining_text = remaining_text[max_length:]
    
    # 过滤空片段
    segments = [seg for seg in segments if seg.strip()]
    return segments

def get_embedding(text, model="bge-m3", max_retries=3):
    """获取文本嵌入向量，带重试机制"""
    for attempt in range(max_retries):
        try:
            response = client.embeddings.create(
                model=model,
                input=[text]
            )
            return response.data[0].embedding
        except Exception as e:
            if attempt < max_retries - 1:
                print(f"⚠️ 获取嵌入失败 (重试 {attempt+1}/{max_retries}): {str(e)}")
                time.sleep(1)
            else:
                print(f"❌ 无法获取文本嵌入: {str(e)}")
                return None
    return None

def enhance_cosine_similarity(query_embed, doc_embed, doc_content):
    """增强版相似度计算：基础相似度 + 关键词权重"""
    if query_embed is None or doc_embed is None:
        return 0.0
    
    a = np.array(query_embed)
    b = np.array(doc_embed)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    if norm_a == 0 or norm_b == 0:
        return 0.0
    
    base_sim = np.dot(a, b) / (norm_a * norm_b)
    
    # 关键词加权
    keyword_boost = 0.0
    for kw in KEYWORDS:
        if kw in doc_content:
            count = doc_content.count(kw)
            keyword_boost += min(0.25, count * 0.1)  # 最大加0.25
    
    return min(base_sim + keyword_boost, 1.0)

def find_top_contexts(query_embedding, documents, top_k=3):
    """查找最相关的上下文片段，优先关键词文档"""
    results = []
    
    for doc in documents:
        if "embedding" not in doc or doc["embedding"] is None:
            continue
            
        sim = enhance_cosine_similarity(query_embedding, doc["embedding"], doc["content"])
        
        # 创建结果对象
        result = {
            "source": doc["source"],
            "content": doc["content"],
            "similarity": sim,
            "has_keywords": doc["has_keywords"]
        }
        
        # 当相似度大于0或包含关键词时考虑
        if sim > 0 or doc["has_keywords"]:
            results.append(result)
    
    if not results:
        return []
    
    # 优先排序：含关键词的文档排前 > 相似度高的排前
    results.sort(key=lambda x: (-x["has_keywords"], -x["similarity"]))
    return results[:top_k]

def build_quality_prompt(user_query, context_results):
    """构建高质量的提示词，过滤低相关片段"""
    prompt = f"用户查询: {user_query}\n\n"
    
    if not context_results:
        return prompt + "⚠️ 未找到与查询相关的上下文信息\n请根据您的知识回答用户查询:"
    
    prompt += "相关上下文信息(基于语义匹配):\n"
    high_relevance = []
    low_relevance = []
    
    for ctx in context_results:
        if ctx["similarity"] >= 0.5 or ctx["has_keywords"]:
            high_relevance.append(ctx)
        else:
            low_relevance.append(ctx)
    
    # 优先添加高相关片段
    for i, ctx in enumerate(high_relevance, 1):
        # 高亮显示关键词
        content = ctx["content"]
        for kw in KEYWORDS:
            if kw in content:
                content = content.replace(kw, f"【{kw}】")
                
        prompt += f"[来源: {ctx['source']}, 相关度: {ctx['similarity']:.1%}]"
        if ctx["has_keywords"]:
            prompt += " 🔑"
        prompt += "\n"
        
        # 添加带缩进的片段内容
        prompt += "> " + "\n> ".join(content.split('\n'))
        prompt += "\n\n"
    
    # 添加低相关片段摘要
    if low_relevance:
        prompt += "---\n其他参考片段:\n"
        for ctx in low_relevance[:2]:
            preview = ctx["content"][:100].replace('\n', ' ') + ("..." if len(ctx["content"]) > 100 else "")
            prompt += f"- {ctx['source']}: {preview}\n"
    
    prompt += "\n请基于以上上下文专业地回答用户查询:"
    return prompt

def calculate_embeddings(self):
    """计算知识嵌入 - 修复输入问题"""
    # 使用文件对话框选择目录
    data_dir = filedialog.askdirectory(title="选择知识库目录")
    if not data_dir:
        return
        
    self.display_message("system", f"🔄 正在计算知识嵌入: {data_dir}...", "system")
    
    # 在后台线程中处理
    def embedding_task():
        try:
            # 加载文档 - 使用原始函数
            documents = load_documents(data_dir)
            
            if not documents:
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    "❌ 没有找到可用的文档内容", 
                    "error"
                ))
                return
            
            self.root.after(0, lambda: self.display_message(
                "system", 
                f"🔄 正在计算 {len(documents)} 个文档的嵌入向量...", 
                "system"
            ))
            
            start_time = time.time()
            for i, doc in enumerate(documents):
                # 使用原始函数获取嵌入
                doc["embedding"] = get_embedding(doc["content"])
                
                # 每10个文档更新一次进度
                if (i+1) % 10 == 0 or i == len(documents)-1:
                    elapsed = time.time() - start_time
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        f"  ⏱️ 已完成 {i+1}/{len(documents)} 个片段 (用时: {elapsed:.1f}s)", 
                        "system"
                    ))
            
            # 保存嵌入结果 - 使用原始逻辑
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            embed_dir = os.path.join(data_dir, "embeddings")
            os.makedirs(embed_dir, exist_ok=True)
            
            for doc in documents:
                if doc["embedding"]:
                    output_path = os.path.join(embed_dir, f"{doc['source']}.embed")
                    with open(output_path, 'w') as f:
                        json.dump({
                            "source": doc["source"],
                            "embedding": doc["embedding"]
                        }, f)
            
            self.root.after(0, lambda: self.display_message(
                "system", 
                f"✅ 已完成知识嵌入计算 | {len(documents)} 个文件 | 保存至: {embed_dir}", 
                "system"
            ))
            
        except Exception as e:
            self.root.after(0, lambda: self.display_message(
                "system", 
                f"❌ 嵌入计算失败: {str(e)}", 
                "error"
            ))
    
    threading.Thread(target=embedding_task, daemon=True).start()
##################################################process#################################################################

def read_all_txt_files(data_dir: str) -> Dict[str, str]:
    """
    从指定目录读取所有 txt 文件的内容
    
    Args:
        data_dir (str): 数据目录路径
        
    Returns:
        Dict[str, str]: 字典，键为文件名，值为文件内容
    """
    txt_files_content = {}
    
    try:
        # 获取数据目录的绝对路径
        data_path = Path(data_dir)
        
        if not data_path.exists():
            print(f"错误：目录 '{data_dir}' 不存在")
            return txt_files_content
        
        if not data_path.is_dir():
            print(f"错误：'{data_dir}' 不是一个目录")
            return txt_files_content
        
        # 查找所有 txt 文件
        txt_files = list(data_path.glob("*.txt"))
        
        if not txt_files:
            print(f"在目录 '{data_dir}' 中没有找到任何 txt 文件")
            return txt_files_content
        
        print(f"在目录 '{data_dir}' 中找到 {len(txt_files)} 个 txt 文件:")
        
        # 读取每个 txt 文件的内容
        for txt_file in txt_files:
            try:
                with open(txt_file, 'r', encoding='utf-8') as f:
                    content = f.read()
                    txt_files_content[txt_file.name] = content
                    print(f"  ✓ 成功读取: {txt_file.name} ({len(content)} 字符)")
            except UnicodeDecodeError:
                # 如果 UTF-8 解码失败，尝试其他编码
                try:
                    with open(txt_file, 'r', encoding='gbk') as f:
                        content = f.read()
                        txt_files_content[txt_file.name] = content
                        print(f"  ✓ 成功读取 (GBK编码): {txt_file.name} ({len(content)} 字符)")
                except Exception as e:
                    print(f"  ✗ 读取文件失败: {txt_file.name} - {e}")
            except Exception as e:
                print(f"  ✗ 读取文件失败: {txt_file.name} - {e}")
        
        return txt_files_content
        
    except Exception as e:
        print(f"读取目录时出错: {e}")
        return txt_files_content

def read_specific_txt_file(file_path: str) -> Optional[str]:
    """
    读取指定的 txt 文件内容
    
    Args:
        file_path (str): 文件路径
        
    Returns:
        Optional[str]: 文件内容，读取失败时返回 None
    """
    try:
        file_path = Path(file_path)
        if not file_path.exists():
            print(f"错误：文件 '{file_path}' 不存在")
            return None
        
        # 获取文件大小
        file_size = os.path.getsize(file_path)
        size_str = f"{file_size} 字节"
        if file_size > 1024:
            size_str = f"{file_size/1024:.2f} KB"
        if file_size > 1024 * 1024:
            size_str = f"{file_size/(1024 * 1024):.2f} MB"
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            print(f"📄 文件: {file_path.name}")
            print(f"📏 大小: {size_str}")
            print(f"📝 字符数: {len(content)}")
            return content
            
    except UnicodeDecodeError:
        try:
            with open(file_path, 'r', encoding='gbk') as f:
                content = f.read()
                print(f"📄 文件: {file_path.name}")
                print(f"📏 大小: {size_str}")
                print(f"📝 字符数: {len(content)}")
                print(f"🔠 编码: GBK")
                return content
        except Exception as e:
            print(f"读取文件失败: {e}")
            return None
    except Exception as e:
        print(f"读取文件失败: {e}")
        return None

def get_txt_files_list(data_dir: str) -> List[str]:
    """
    获取指定目录中所有 txt 文件的列表
    
    Args:
        data_dir (str): 数据目录路径
        
    Returns:
        List[str]: txt 文件名列表
    """
    try:
        data_path = Path(data_dir)
        if not data_path.exists():
            print(f"错误：目录 '{data_dir}' 不存在")
            return []
        
        txt_files = [f.name for f in data_path.glob("*.txt")]
        return txt_files
        
    except Exception as e:
        print(f"获取文件列表时出错: {e}")
        return []

def process():
    """主处理函数"""
    print("=" * 60)
    print("📁 TXT 文件读取工具")
    print("=" * 60)
    
    # 获取用户输入的目录名称
    data_dir = input("请输入要读取的目录名称: ").strip()
    
    # 主循环
    while True:
        print("\n" + "=" * 60)
        print("🔍 主菜单")
        print("=" * 60)
        print("1. 查看文件列表")
        print("2. 读取特定文件")
        print("3. 读取所有文件")
        print("4. 退出程序")
        
        choice = input("\n请选择操作 (1/2/3/4): ").strip()
        
        if choice == "1":
            # 获取文件列表
            file_list = get_txt_files_list(data_dir)
            if file_list:
                print(f"\n在 '{data_dir}' 中找到 {len(file_list)} 个 txt 文件:")
                for file in file_list:
                    print(f"  - {file}")
            else:
                print(f"在 '{data_dir}' 中没有找到 txt 文件")
        
        elif choice == "2":
            # 读取特定文件
            file_name = input("\n请输入要读取的文件名: ").strip()
            if not file_name:
                print("文件名不能为空")
                continue
            
            file_path = os.path.join(data_dir, file_name)
            content = read_specific_txt_file(file_path)
            
            if content:
                print("\n" + "=" * 60)
                print("📄 文件内容预览:")
                print("=" * 60)
                
                # 显示前500个字符作为预览
                preview = content[:500]
                print(preview)
                
                if len(content) > 500:
                    print(f"... (共 {len(content)} 字符)")
                
                # 提供查看完整内容的选项
                view_full = input("\n是否查看完整内容? (y/n): ").strip().lower()
                if view_full == 'y':
                    print("\n" + "=" * 60)
                    print("📄 完整文件内容:")
                    print("=" * 60)
                    print(content)
                    print("=" * 60)
        
        elif choice == "3":
            # 读取所有文件
            print(f"\n正在读取 '{data_dir}' 目录中的所有 txt 文件...")
            all_content = read_all_txt_files(data_dir)
            
            if all_content:
                print("\n" + "=" * 60)
                print("📊 文件摘要")
                print("=" * 60)
                print(f"共读取 {len(all_content)} 个文件:")
                
                for filename, content in all_content.items():
                    print(f"\n📄 文件: {filename}")
                    print(f"📏 字符数: {len(content)}")
                    
                    # 显示前100个字符作为预览
                    preview = content[:100].replace('\n', '\\n')
                    if len(content) > 100:
                        preview += "..."
                    print(f"📝 内容预览: {preview}")
            else:
                print("没有读取到任何文件内容")
        
        elif choice == "4":
            print("\n感谢使用，程序结束!")
            break
        
        else:
            print("\n⚠️ 无效选择，请重新输入")
        
        # 添加分隔线
        print("\n" + "-" * 60)
def create_client():
    """创建OpenAI客户端 - 完整实现"""
    try:
        # 使用安全的httpx客户端
        import httpx
        http_client = httpx.Client(
            timeout=60.0,
            limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
            transport=httpx.HTTPTransport(retries=3)
        )
        
        from openai import OpenAI
        return OpenAI(
            api_key='sk-uZy29fuiHwjQEHtJq35r5KQHDqYYMa3WNdVdvxo8ZoPyws4T',
            base_url='http://123.182.124.194:58000/v1',
            http_client=http_client
        )
    except Exception as e:
        print(f"创建客户端时出错: {e}")
        return None

def stream_chat(client, messages):
    """流式聊天函数 - 完整实现"""
    try:
        # 创建流式响应
        response = client.chat.completions.create(
            model="Qwen2.5-72B-Instruct",
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=2000
        )
        
        full_response = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                full_response += content
        return full_response
        
    except Exception as e:
        return f"API调用错误: {str(e)}"

def build_rag_prompt(user_input, knowledge, history):
    """构建RAG提示词 - 完整实现"""
    system_msg = {
        "role": "system",
        "content": "你是一个友好且专业的AI助手，请用中文回答用户问题"
    }
    
    # 选择与问题最相关的知识片段
    keyword_hits = []
    for item in knowledge:
        path_content = f"{item['path']}\n{item['content']}"
        if any(keyword in path_content for keyword in KEYWORDS):
            keyword_hits.append(path_content)
    
    # 优先包含包含关键字的片段
    knowledge_context = ""
    for hit in keyword_hits[:5]:
        knowledge_context += f"【相关片段】\n{hit}\n\n"
    
    # 添加其他知识片段
    count = 0
    for item in knowledge:
        path_content = f"{item['path']}\n{item['content']}"
        if path_content not in keyword_hits:  # 避免重复
            knowledge_context += f"【{item['path']}】\n{item['content'][:500]}...\n\n"
            count += 1
            if count >= 5:  # 最多显示5个额外片段
                break
    
    if not knowledge_context:
        knowledge_context = "⚠️ 未加载相关背景知识"
    
    knowledge_msg = {
        "role": "user",
        "content": f"{knowledge_context}\n\n【当前问题】{user_input}"
    }
    
    return [system_msg] + history + [knowledge_msg]

def terminate_process():
    """强制终止当前进程和所有子进程"""
    try:
        if sys.platform == "win32":
            # Windows 系统
            import ctypes
            ctypes.windll.kernel32.ExitProcess(0)
        else:
            # Unix/Linux/Mac 系统
            os.kill(os.getpid(), signal.SIGTERM)
    except Exception as e:
        print(f"无法终止进程: {e}")
        sys.exit(1)


# ======================== GUI封装层 - 完整功能 ========================
class TerminalStyleChat:
    def __init__(self, root):
        self.root = root
        self.root.title("AI助手 - 终端风格聊天系统")
        self.root.geometry("900x700")
        self.root.configure(bg="#1e1e1e")
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.api_key = "a71db04526394b7b85ce4517341c198e.FnqpGsOJwLW7sCrK"
        try:
            self.ocr_processor = ImageOCR(self.api_key)
            print("✅ OCR处理器初始化成功")
        except Exception as ex:
            print(f"❌ OCR处理器初始化失败: {str(ex)}")
            self.ocr_processor = None
        
        # 初始化功能模块
        self.client = create_client()
        if not self.client:
            messagebox.showerror("错误", "无法创建客户端，程序退出")
            return
        
        # 初始化管理器
        self.crawler = WebCrawler()
        self.history_manager = HistoryManager()
        self.knowledge_loader = KnowledgeLoader()
        
        # 初始化OCR处理器
        self.api_key = os.getenv("GLM_API_KEY")
        self.ocr_processor = ImageOCR(self.api_key) if self.api_key else None
        
        # 加载知识库
        self.knowledge = self.knowledge_loader.load()
        
        # 初始化对话
        self.messages = [
            {"role": "system", "content": "你是一个友好且专业的AI助手，请用中文回答用户问题"}
        ]
        
        # 尝试加载最新记录
        self.latest_history = self.history_manager.load_history("latest.json")
        if self.latest_history:
            self.messages = self.latest_history
        
        # 创建UI
        self.create_widgets()
        
        # 显示初始消息
        self.display_message("system", "欢迎使用终端风格聊天系统! 输入 'help' 查看可用命令。", "system")
    
    def create_widgets(self):
        """创建UI组件"""
        # 主框架
        self.main_frame = Frame(self.root, bg="#1e1e1e")
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 标题区域
        self.create_header()
        
        # 聊天显示区域
        self.create_chat_display()
        
        # 输入区域
        self.create_input_area()
        
        # 状态栏
        self.create_status_bar()
    
    def create_header(self):
        """创建标题区域"""
        header_frame = Frame(self.main_frame, bg="#2d2d2d", height=50)
        header_frame.pack(fill=tk.X, pady=(0, 10))
        
        title_label = Label(
            header_frame, 
            text="🤖 AI助手 - 终端风格聊天系统", 
            fg="#4ec9b0",
            bg="#2d2d2d",
            font=("Consolas", 14, "bold")
        )
        title_label.pack(side=tk.LEFT, padx=15)
        
        # 关闭按钮
        close_button = Button(
            header_frame,
            text="退出",
            bg="#f44747",
            fg="white",
            font=("Consolas", 10, "bold"),
            relief=tk.FLAT,
            command=self.force_exit
        )
        close_button.pack(side=tk.RIGHT, padx=15)
        
        # 状态指示灯
        self.status_indicator = Label(
            header_frame, 
            text="●", 
            fg="#4ec9b0",
            bg="#2d2d2d",
            font=("Arial", 16)
        )
        self.status_indicator.pack(side=tk.RIGHT, padx=15)
    
    def create_chat_display(self):
        """创建聊天显示区域"""
        chat_frame = Frame(self.main_frame, bg="#1e1e1e")
        chat_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建滚动文本框
        self.chat_display = scrolledtext.ScrolledText(
            chat_frame,
            wrap=tk.WORD,
            bg="#1e1e1e",
            fg="#d4d4d4",
            insertbackground="#d4d4d4",
            font=("Consolas", 12),
            padx=15,
            pady=15,
            state=tk.DISABLED
        )
        self.chat_display.pack(fill=tk.BOTH, expand=True)
        
        # 配置标签样式
        self.chat_display.tag_configure("user", foreground="#569cd6")
        self.chat_display.tag_configure("assistant", foreground="#4ec9b0")
        self.chat_display.tag_configure("system", foreground="#c586c0")
        self.chat_display.tag_configure("command", foreground="#dcdcaa")
        self.chat_display.tag_configure("error", foreground="#f44747")
    
    def create_input_area(self):
        """创建输入区域"""
        input_frame = Frame(self.main_frame, bg="#1e1e1e")
        input_frame.pack(fill=tk.X, pady=(10, 0))
        
        # 输入提示符
        prompt_label = Label(
            input_frame, 
            text=">>>", 
            fg="#4ec9b0",
            bg="#1e1e1e",
            font=("Consolas", 12)
        )
        prompt_label.pack(side=tk.LEFT)
        
        # 输入框
        self.user_input = Entry(
            input_frame,
            bg="#252526",
            fg="#d4d4d4",
            insertbackground="#d4d4d4",
            font=("Consolas", 12),
            relief=tk.FLAT,
            width=80
        )
        self.user_input.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        self.user_input.bind("<Return>", self.send_message)
        self.user_input.focus_set()
        
        # 发送按钮
        send_button = Button(
            input_frame,
            text="发送",
            bg="#007acc",
            fg="white",
            font=("Consolas", 10, "bold"),
            relief=tk.FLAT,
            command=self.send_message
        )
        send_button.pack(side=tk.RIGHT)
    
    def create_status_bar(self):
        """创建状态栏"""
        status_frame = Frame(self.main_frame, bg="#007acc", height=25)
        status_frame.pack(fill=tk.X, pady=(10, 0))
        
        # 左侧状态信息
        self.status_var = tk.StringVar(value=f"就绪 | 知识库: {len(self.knowledge)} 个文件")
        status_label = Label(
            status_frame,
            textvariable=self.status_var,
            fg="white",
            bg="#007acc",
            font=("Consolas", 10)
        )
        status_label.pack(side=tk.LEFT, padx=10)
        
        # 右侧时间信息
        self.time_var = tk.StringVar(value=get_current_datetime())
        time_label = Label(
            status_frame,
            textvariable=self.time_var,
            fg="white",
            bg="#007acc",
            font=("Consolas", 10)
        )
        time_label.pack(side=tk.RIGHT, padx=10)
        
        # 更新时间
        def update_time():
            self.time_var.set(get_current_datetime())
            self.root.after(1000, update_time)
        
        update_time()
    
    def update_status(self, message):
        """更新状态栏信息"""
        self.status_var.set(message)
    
    def display_message(self, sender, message, tag="system"):
        """在聊天区域显示消息"""
        self.chat_display.config(state=tk.NORMAL)
        
        # 添加发送者标签
        if sender == "user":
            prefix = "👤 您: "
        elif sender == "assistant":
            prefix = "🤖 AI助手: "
        else:
            prefix = "🖥️ 系统: "
        
        self.chat_display.insert(tk.END, prefix, tag)
        
        # 添加消息内容
        self.chat_display.insert(tk.END, message + "\n\n", tag)
        
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.yview(tk.END)  # 滚动到底部
    
    def send_message(self, event=None):
        """发送用户消息"""
        user_text = self.user_input.get().strip()
        if not user_text:
            return
        
        # 清空输入框
        self.user_input.delete(0, tk.END)
        
        # 显示用户消息
        self.display_message("user", user_text, "user")
        
        # 处理命令
        if self.handle_command(user_text):
            return
        
        # 在独立线程中处理AI响应
        threading.Thread(target=self.get_ai_response, args=(user_text,), daemon=True).start()
    
    def handle_command(self, command):
        """处理特殊命令"""
        command_lower = command.lower()
        
        # 命令映射表
        command_map = {
            "1": "list", "list": "list",
            "2": "load", "load": "load",
            "3": "save", "save": "save",
            "4": "clear", "clear": "clear",
            "5": "reload", "reload": "reload",
            "6": "kload", "kload": "kload",
            "7": "exit", "exit": "exit",
            "8": "embedding", "embedding": "embedding",
            "9": "process", "process": "process",
            "10": "crawl", "crawl": "crawl",
            "11": "ocr", "ocr": "ocr",
            "help": "help"
        }
        
        # 处理命令映射
        if command_lower in command_map:
            command = command_map[command_lower]
        
        # 命令处理
        if command == 'help':
            self.show_help()
            return True
            
        if command == 'exit':
            self.save_conversation()
            self.root.destroy()
            return True
            
        if command == 'list':
            self.list_histories()
            return True
            
        if command_lower.startswith('load '):
            self.load_histories(command_lower[5:])
            return True
            
        if command_lower.startswith('kload '):
            self.load_knowledge(command_lower[6:])
            return True
            
        if command == 'save':
            self.save_conversation()
            return True
            
        if command == 'clear':
            self.clear_conversation()
            return True
            
        if command == 'reload':
            self.reload_knowledge()
            return True
            
        if command == 'embedding':
            self.calculate_embeddings()
            return True
            
        if command == 'process':
            self.process_files()
            return True
            
        if command_lower.startswith('crawl '):
            self.crawl_website(command_lower[6:])
            return True
            
        if command == 'ocr':
            self.process_ocr()
            return True
            
        return False
    
    def show_help(self):
        """显示帮助信息"""
        help_text = """
可用命令:
  help    - 显示帮助信息
  exit    - 退出程序
  list    - 列出历史对话
  load <编号1>,<编号2> - 加载多个历史文件
  save    - 保存当前对话
  clear   - 清空当前对话
  reload  - 重载知识库
  kload <文件模式> - 选择知识库
  embedding - 计算知识库嵌入
  process - 查看目录文件
  crawl <URL> [dynamic] [selector] - 爬取网页内容
  ocr     - 识别图片文字
"""
        self.display_message("system", help_text.strip(), "system")
    
    def list_histories(self):
        """列出历史对话文件"""
        files = self.history_manager.list_histories()
        if not files:
            self.display_message("system", "ℹ️ 无历史记录", "system")
            return
        
        self.display_message("system", "📜 历史对话记录:", "system")
        for file in files:
            self.display_message("system", f"  [{file['id']}] {file['filename']}", "system")
    
    def load_histories(self, command_arg=""):
        """加载多个历史文件"""
        if not command_arg:
            self.display_message("system", "请输入要加载的文件编号，例如: load 1,3,5", "command")
            return
        
        try:
            # 解析多个ID
            ids = [int(id.strip()) for id in command_arg.split(',')]
            files = self.history_manager.list_histories()
            
            # 获取文件名
            filenames = []
            for id in ids:
                if 1 <= id <= len(files):
                    filenames.append(files[id-1]["filename"])
            
            if filenames:
                # 加载并合并多个历史文件
                merged_history = self.history_manager.load_multiple_histories(filenames)
                self.messages = merged_history
                self.update_chat_display()
                self.display_message("system", f"✅ 已加载 {len(filenames)} 个历史文件", "system")
            else:
                self.display_message("system", "❌ 未找到匹配的历史文件", "error")
        except:
            self.display_message("system", "⚠️ 命令格式: load <编号1>,<编号2>,...", "error")
    
    def update_chat_display(self):
        """更新聊天历史显示"""
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.delete(1.0, tk.END)
        
        for msg in self.messages:
            if msg["role"] == "system":
                continue
                
            prefix = "👤 您: " if msg["role"] == "user" else "🤖 AI助手: "
            tag = "user" if msg["role"] == "user" else "assistant"
            
            self.chat_display.insert(tk.END, prefix, tag)
            self.chat_display.insert(tk.END, msg["content"] + "\n\n")
        
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.yview(tk.END)
    
    def load_knowledge(self, command_arg=""):
        """加载知识库"""
        if command_arg:
            if command_arg.startswith('content:'):
                keyword = command_arg[8:]
                def content_filter(text):
                    return keyword.lower() in text.lower()
                self.knowledge = self.knowledge_loader.kload(content_filter=content_filter)
                self.update_status(f"知识库: {len(self.knowledge)} 个文件")
                self.display_message("system", f"🔍 已加载包含关键词 '{keyword}' 的 {len(self.knowledge)} 个文件", "system")
            else:
                pattern = command_arg
                self.knowledge = self.knowledge_loader.kload(file_filter=pattern)
                self.update_status(f"知识库: {len(self.knowledge)} 个文件")
                self.display_message("system", f"📂 已按模式 '{pattern}' 加载 {len(self.knowledge)} 个文件", "system")
        else:
            # 打开文件对话框选择知识库文件
            file_path = filedialog.askopenfilename(
                title="选择知识库文件",
                filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
            )
            if file_path:
                # 这里应包含加载特定知识库文件的逻辑
                self.display_message("system", f"✅ 已加载知识库文件: {os.path.basename(file_path)}", "system")
    
    def save_conversation(self):
        """保存当前对话"""
        saved_path = self.history_manager.save_history(self.messages)
        if saved_path:
            self.display_message("system", f"💾 对话已保存至: {saved_path}", "system")
        else:
            self.display_message("system", "❌ 保存对话失败", "error")
    
    def clear_conversation(self):
        """清空当前对话"""
        if messagebox.askyesno("确认", "确定要清空当前对话吗?"):
            self.messages = [self.messages[0]]
            self.update_chat_display()
            self.display_message("system", "🗑️ 对话已清空", "system")
    
    def reload_knowledge(self):
        """重载知识库"""
        self.knowledge = self.knowledge_loader.load()
        self.update_status(f"就绪 | 知识库: {len(self.knowledge)} 个文件")
        self.display_message("system", f"🔄 已重载 {len(self.knowledge)} 个知识文件", "system")
    
    def calculate_embeddings(self):
        """计算知识嵌入 - 修复输入问题"""
        # 使用文件对话框选择目录
        data_dir = filedialog.askdirectory(title="选择知识库目录")
        if not data_dir:
            return
            
        self.display_message("system", f"🔄 正在计算知识嵌入: {data_dir}...", "system")
        
        # 在后台线程中处理
        def embedding_task():
            try:
                # 加载文档 - 使用原始函数
                documents = load_documents(data_dir)
                
                if not documents:
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        "❌ 没有找到可用的文档内容", 
                        "error"
                    ))
                    return
                
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"🔄 正在计算 {len(documents)} 个文档的嵌入向量...", 
                    "system"
                ))
                
                start_time = time.time()
                for i, doc in enumerate(documents):
                    # 使用原始函数获取嵌入
                    doc["embedding"] = get_embedding(doc["content"])
                    
                    # 每10个文档更新一次进度
                    if (i+1) % 10 == 0 or i == len(documents)-1:
                        elapsed = time.time() - start_time
                        self.root.after(0, lambda: self.display_message(
                            "system", 
                            f"  ⏱️ 已完成 {i+1}/{len(documents)} 个片段 (用时: {elapsed:.1f}s)", 
                            "system"
                        ))
                
                # 保存嵌入结果 - 使用原始逻辑
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                embed_dir = os.path.join(data_dir, "embeddings")
                os.makedirs(embed_dir, exist_ok=True)
                
                for doc in documents:
                    if doc["embedding"]:
                        output_path = os.path.join(embed_dir, f"{doc['source']}.embed")
                        with open(output_path, 'w') as f:
                            json.dump({
                                "source": doc["source"],
                                "embedding": doc["embedding"]
                            }, f)
                
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"✅ 已完成知识嵌入计算 | {len(documents)} 个文件 | 保存至: {embed_dir}", 
                    "system"
                ))
                
            except Exception as e:
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"❌ 嵌入计算失败: {str(e)}", 
                    "error"
                ))
        
        threading.Thread(target=embedding_task, daemon=True).start()
    
    def process_files(self):
        """处理目录文件 - 完全恢复原始功能"""
        # 选择要处理的目录
        data_dir = filedialog.askdirectory(title="选择要处理的目录")
        if not data_dir:
            return
            
        self.display_message("system", f"🔄 正在处理目录: {data_dir}...", "system")
        
        # 在后台线程中处理
        def process_task():
            try:
                # 获取所有文件 - 使用原始逻辑
                files = glob.glob(os.path.join(data_dir, '*.txt'))
                total = len(files)
                
                if not files:
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        f"❌ 在 '{data_dir}' 中没有找到任何txt文件", 
                        "error"
                    ))
                    return
                
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"📂 发现 {total} 个文本文件", 
                    "system"
                ))
                
                result = []
                for i, file_path in enumerate(files):
                    try:
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            content = f.read()
                            result.append({
                                "path": os.path.basename(file_path),
                                "size": os.path.getsize(file_path),
                                "chars": len(content)
                            })
                    except:
                        pass
                    
                    # 更新进度
                    if (i+1) % 10 == 0 or i == total - 1:
                        self.root.after(0, lambda: self.display_message(
                            "system", 
                            f"处理中... {i+1}/{total} 个文件", 
                            "system"
                        ))
                
                # 生成报告 - 使用原始格式
                report = "\n".join([
                    f"{r['path']} - 大小: {r['size']:,} 字节, 字符数: {r['chars']:,}"
                    for r in result
                ])
                
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"✅ 已处理 {total} 个文件:\n" + report, 
                    "system"
                ))
                
            except Exception as e:
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"❌ 处理失败: {str(e)}", 
                    "error"
                ))
        
        threading.Thread(target=process_task, daemon=True).start()
    
    def get_ai_response(self, user_input):
        """获取AI响应 - 完全恢复原始功能"""
        # 添加用户消息到历史
        self.messages.append({"role": "user", "content": user_input})
        
        # 在独立线程中处理API调用
        def api_task():
            try:
                # 构建RAG提示 - 使用原始函数
                rag_prompt = build_rag_prompt(user_input, self.knowledge, self.messages)
                
                # 使用原始流式聊天函数
                response = stream_chat(self.client, rag_prompt)
                
                if response.startswith("API调用错误"):
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        response, 
                        "error"
                    ))
                    # 移除无效的用户消息
                    self.messages.pop()
                else:
                    # 添加AI回复到历史
                    self.messages.append({"role": "assistant", "content": response})
                    
                    # 更新UI - 流式显示
                    lines = response.split('\n')
                    for i, line in enumerate(lines):
                        self.root.after(0, lambda l=line: self.display_message_partial("assistant", l))
                        time.sleep(0.05)
                        
                        # 添加换行
                        if i < len(lines) - 1:
                            self.root.after(0, lambda: self.display_message_partial("assistant", "\n"))
                    
                    # 添加消息分隔符
                    self.root.after(0, lambda: self.display_message_partial("assistant", "\n\n"))
                    
                    # 保持上下文长度
                    if len(self.messages) > 50:
                        self.messages = [self.messages[0]] + self.messages[-49:]
            except Exception as e:
                error_msg = f"❌ 获取AI响应时出错: {str(e)}"
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    error_msg, 
                    "error"
                ))
                self.messages.pop()
        
        threading.Thread(target=api_task, daemon=True).start()
    
    def display_message_partial(self, sender, message, tag="assistant"):
        """部分显示消息 - 用于流式输出"""
        self.chat_display.config(state=tk.NORMAL)
        
        if self.chat_display.get("end-1c", "end") != "\n":
            self.chat_display.insert(tk.END, " ")
            
        # 添加部分消息
        if message:
            self.chat_display.insert(tk.END, message, tag)
        
        self.chat_display.config(state=tk.DISABLED)
        self.chat_display.see(tk.END)  # 滚动到底部
    
    def crawl_website(self, command_arg=""):
        """爬取网页内容"""
        # 解析参数
        args = command_arg.split() if command_arg else []
        
        # 创建爬取对话框
        crawl_window = tk.Toplevel(self.root)
        crawl_window.title("网页爬取")
        crawl_window.geometry("500x300")
        crawl_window.configure(bg="#2d2d2d")
        
        ttk.Label(crawl_window, text="URL:", bg="#2d2d2d", fg="#d4d4d4").pack(anchor=tk.W, padx=10, pady=(10, 0))
        url_entry = Entry(crawl_window, bg="#252526", fg="#d4d4d4", font=("Consolas", 10))
        url_entry.pack(fill=tk.X, padx=10, pady=5)
        
        if args:
            url_entry.insert(0, args[0])
        
        # 选项
        options_frame = Frame(crawl_window, bg="#2d2d2d")
        options_frame.pack(fill=tk.X, padx=10, pady=5)
        
        dynamic_var = tk.BooleanVar(value='dynamic' in args[1:])
        dynamic_check = tk.Checkbutton(
            options_frame, 
            text="动态页面", 
            variable=dynamic_var,
            bg="#2d2d2d",
            fg="#d4d4d4",
            selectcolor="#1e1e1e"
        )
        dynamic_check.pack(side=tk.LEFT, padx=(0, 10))
        
        tk.Label(options_frame, text="选择器:", bg="#2d2d2d", fg="#d4d4d4").pack(side=tk.LEFT)
        selector_entry = Entry(options_frame, bg="#252526", fg="#d4d4d4", font=("Consolas", 10), width=20)
        selector_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 添加默认选择器（如果有）
        selector = next((arg for arg in args[1:] if arg not in ('dynamic',)), None)
        if selector:
            selector_entry.insert(0, selector)
        
        # 按钮
        button_frame = Frame(crawl_window, bg="#2d2d2d")
        button_frame.pack(fill=tk.X, padx=10, pady=10)
        
        def start_crawl():
            url = url_entry.get().strip()
            if not url:
                self.display_message("system", "❌ 请输入URL", "error")
                return
                
            is_dynamic = dynamic_var.get()
            selector = selector_entry.get().strip() or None
            
            # 执行爬取
            self.display_message("system", f"🔄 正在爬取: {url}...", "system")
            
            def crawl_task():
                try:
                    result = self.crawler.crawl_page(url, is_dynamic, selector)
                    
                    if result.get('status') == 'success':
                        self.root.after(0, lambda: self.display_message(
                            "system", 
                            f"✅ 爬取成功!\n📁 保存路径: {result.get('path')}\n📝 内容预览: {result.get('content')[:500]}...", 
                            "system"
                        ))
                        crawl_window.destroy()
                    else:
                        self.root.after(0, lambda: self.display_message(
                            "system", 
                            f"❌ 爬取失败: {result.get('reason')}", 
                            "error"
                        ))
                except Exception as e:
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        f"❌ 爬取失败: {str(e)}", 
                        "error"
                    ))
            
            threading.Thread(target=crawl_task, daemon=True).start()
        
        tk.Button(
            button_frame,
            text="开始爬取",
            bg="#007acc",
            fg="white",
            font=("Consolas", 10, "bold"),
            command=start_crawl
        ).pack(side=tk.LEFT)
        
        tk.Button(
            button_frame,
            text="取消",
            bg="#f44747",
            fg="white",
            font=("Consolas", 10, "bold"),
            command=crawl_window.destroy
        ).pack(side=tk.RIGHT)
    
    def process_ocr(self):
        """处理OCR图片识别 - 添加处理器检查"""
        # 检查OCR处理器是否可用
        if self.ocr_processor is None:
            self.display_message("system", "❌ OCR功能不可用，请检查API密钥", "error")
            return
            
        # 打开文件对话框选择图片
        file_path = filedialog.askopenfilename(
            title="选择图片文件",
            filetypes=[("图片文件", "*.png *.jpg *.jpeg *.bmp")]
        )
        
        if not file_path:
            return
            
        self.display_message("system", f"🔄 正在识别图片: {os.path.basename(file_path)}...", "system")
        
        # 在后台线程中处理OCR
        threading.Thread(
            target=self.perform_ocr, 
            args=(file_path,),
            daemon=True
        ).start()
    def get_api_key(self):
        """获取 API 密钥 - 优先使用环境变量"""
        # 1. 尝试从环境变量获取
        api_key = os.getenv("GLM_API_KEY")
        if api_key:
            return api_key
        
        # 2. 尝试从配置文件获取
        try:
            with open("api_key.txt", "r") as f:
                return f.read().strip()
        except:
            pass
        
        # 3. 使用默认密钥
        default_key = "a71db04526394b7b85ce4517341c198e.FnqpGsOJwLW7sCrK"
        from tkinter import simpledialog
        # 4. 提示用户输入
        user_key = simpledialog.askstring(
            "API密钥输入",
            "请输入GLM_API_KEY (留空使用默认密钥):",
            parent=self.root
        )
        
        return user_key.strip() if user_key else default_key
    def perform_ocr(self, file_path):
        """执行OCR识别 - 添加处理器检查"""
        # 再次检查OCR处理器是否可用
        if self.ocr_processor is None:
            self.root.after(0, lambda: self.display_message(
                "system", 
                "❌ OCR功能不可用，无法执行识别", 
                "error"
            ))
            return
            
        try:
            # 执行OCR识别
            text_content = self.ocr_processor.recognize_image(file_path)
            
            if text_content:
                # 显示结果
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    f"✅ OCR识别完成:\n文件: {os.path.basename(file_path)}\n内容: {text_content[:500]}...", 
                    "system"
                ))
                
                # 保存结果
                output_path = self.ocr_processor.save_ocr_result(file_path)
                if output_path:
                    self.root.after(0, lambda: self.display_message(
                        "system", 
                        f"💾 OCR结果已保存至: {output_path}", 
                        "system"
                    ))
            else:
                self.root.after(0, lambda: self.display_message(
                    "system", 
                    "❌ OCR识别失败", 
                    "error"
                ))
        except Exception as ex:
            self.root.after(0, lambda ex=ex: self.display_message(
                "system", 
                f"❌ OCR处理出错: {str(ex)}", 
                "error"
            ))
    
    def force_exit(self):
        """强制退出应用程序"""
        if messagebox.askyesno("强制退出", "确定要强制退出应用程序吗? 所有未保存的数据将会丢失!"):
            self.root.destroy()
    
    def on_close(self):
        """窗口关闭时的处理"""
        if messagebox.askyesno("确认退出", "确定要退出程序吗?"):
            try:
                saved_path = self.history_manager.save_history(self.messages)
                if saved_path:
                    self.display_message("system", f"💾 对话已保存至: {saved_path}", "system")
            except Exception as e:
                traceback.print_exc()
                messagebox.showerror("保存错误", f"保存对话时出错: {str(e)}\n程序将强制退出")
            finally:
                self.root.after(500, self.root.destroy)

# ======================== 主程序入口 ========================
if __name__ == "__main__":
    root = tk.Tk()
    app = TerminalStyleChat(root)
    root.mainloop()